// @HEADER
//*********************************************************************//
//  SiYuan: A numerical PDE solver                                     //
//  Copyright (2022) YUAN Xi                                           //
//  This Software is released under the BSD 2-Clause license detailed  //
//  in the file "LICENSE" in the top-level SiYuan directory            //
//*********************************************************************//
// @HEADER

#ifndef _EquationSet_Magnetostatics_hpp_
#define _EquationSet_Magnetostatics_hpp_

#include <vector>
#include <string>

// Kokkos
#include "Kokkos_DynRankView.hpp"

#include "Teuchos_RCP.hpp"
#include "Panzer_EquationSet_DefaultImpl.hpp"
#include "Panzer_Traits.hpp"
#include "Phalanx_FieldManager.hpp"

namespace SiYuan {

  template <typename EvalT>
  class EquationSet_Magnetostatics : public EquationSet<EvalT> {

  public:    

    EquationSet_Magnetostatics(const Teuchos::RCP<Teuchos::ParameterList>& params,
		       const int& default_integration_order,
		       const panzer::CellData& cell_data,
		       const Teuchos::RCP<panzer::GlobalData>& gd,
		       const bool build_transient_support);
    
    void buildAndRegisterEquationSetEvaluators(PHX::FieldManager<panzer::Traits>& fm,
					const panzer::FieldLibrary& field_library,
					const Teuchos::ParameterList& user_data) const;

  private:
      std::string m_dof_name;
      std::string m_matl_name;
      bool has_M, has_J;

      Teuchos::Array<double> m_M;
      Teuchos::Array<double> m_J;
  };

  template<typename EvalT, typename Traits>
  class Residual_Magnetostatics : public panzer::EvaluatorWithBaseImpl<Traits>,
    public PHX::EvaluatorDerived<EvalT, Traits>
  {
    public:

      /**
       *  \brief Main Constructor.
       *
       *  Creates an `Evaluator` to evaluate the integral
       *  \f[
       *    Ma(x)b(x)\cdots\int\vec{s}(x)\cdot\nabla\phi(x)\,dx,
       *  \f]
       *  where \f$ M \f$ is some constant, \f$ a(x) \f$, \f$ b(x) \f$, etc.,
       *  are some fields that depend on position, \f$ \vec{s} \f$ is some
       *  vector-valued function, and \f$ \phi \f$ is some basis.
       *
       *  \param[in] evalStyle  An `enum` declaring the behavior of this
       *                        `Evaluator`, which is to either:
       *                        - compute and contribute (`CONTRIBUTES`), or
       *                        - compute and store (`EVALUATES`).
       *  \param[in] resName    The name of either the contributed or evaluated
       *                        field, depending on `evalStyle`.
       *  \param[in] fluxName   The name of the vector-valued function being
       *                        integrated (\f$ \vec{s} \f$).
       *  \param[in] basis      The basis that you'd like to use (\f$ \phi
       *                        \f$).
       *  \param[in] ir         The integration rule that you'd like to use.
       *  \param[in] vecDL      The vector data layout that you'd like to use.
       *                        If not specified, this defaults to
       *                        `Teuchos::null` and the vector data layout from
       *                        the given `ir` is used.
       *
       *  \throws std::invalid_argument If any of the inputs are invalid.
       *  \throws std::logic_error      If the `basis` supplied does not
       *                                support the gradient operator, or if
       *                                the dimension of the space exceeds the
       *                                dimension of the vector data layout.
       */
      Residual_Magnetostatics(
        const panzer::EvaluatorStyle&   evalStyle,
        const std::string&              resName,
        const std::string&              valName,
        const panzer::BasisIRLayout&    basis,
        const panzer::IntegrationRule&  ir );

      /**
       *  \brief Post-Registration Setup.
       *
       *  Sets the basis index, and gets the Kokkos::View versions of the field
       *  multiplier, if there are any.
       *
       *  \param[in] sd Essentially a list of `Workset`s, which are collections
       *                of cells (elements) that all live on a single process.
       *  \param[in] fm This is an unused part of the `Evaluator` interface.
       */
      void postRegistrationSetup( typename Traits::SetupData d,
        PHX::FieldManager<Traits>& fm);

      /**
       *  \brief Evaluate Fields.
       *
       *  This actually performs the integration by calling `operator()()` in a
       *  `Kokkos::parallel_for` over the cells in the `Workset`.
       *
       *  \param[in] workset The `Workset` on which you're going to do the
       *                     integration.
       */
      void evaluateFields(typename Traits::EvalData d);

      /**
       *  \brief This empty struct allows us to optimize
       *         `operator()()` depending on the number of field
       *         multipliers. This is the version that does not use
       *         shared memory.
       */
      struct FieldMultTag {};

      /**
       *  \brief This empty struct allows us to optimize
       *         `operator()()` depending on the number of field
       *         multipliers. This is the shared memory version.
       */
      struct SharedFieldMultTag {};

      /**
       *  \brief Perform the integration.
       *
       *  Generally speaking, for a given cell in the `Workset`, this routine
       *  loops over quadrature points, vector dimensions, and bases to perform
       *  the integration, scaling the vector field to be integrated by the
       *  multiplier (\f$ M \f$) and any field multipliers (\f$ a(x) \f$, \f$
       *  b(x) \f$, etc.).
       *
       *  \note Optimizations are made for the cases in which we have no field
       *        multipliers or only a single one.
       *
       *  \param[in] tag  An indication of the number of field multipliers we
       *                  have; either 0, 1, or something else.
       *  \param[in] cell The cell in the `Workset` over which to integrate.
       */
      KOKKOS_INLINE_FUNCTION
      void
      operator()(
        const FieldMultTag& tag,
        const Kokkos::TeamPolicy<PHX::exec_space>::member_type& team) const;

      /**
       *  \brief Perform the integration.
       *
       *  Generally speaking, for a given cell in the `Workset`, this routine
       *  loops over quadrature points, vector dimensions, and bases to perform
       *  the integration, scaling the vector field to be integrated by the
       *  multiplier (\f$ M \f$) and any field multipliers (\f$ a(x) \f$, \f$
       *  b(x) \f$, etc.). Uses Shared memory.
       *
       *  \note Optimizations are made for the cases in which we have no field
       *        multipliers or only a single one.
       *
       *  \param[in] tag  An indication of the number of field multipliers we
       *                  have; either 0, 1, or something else.
       *  \param[in] cell The cell in the `Workset` over which to integrate.
       */
      KOKKOS_INLINE_FUNCTION
      void
      operator()(
        const SharedFieldMultTag& tag,
        const Kokkos::TeamPolicy<PHX::exec_space>::member_type& team) const;

    private:

      /**
       *  \brief The scalar type.
       */
      using ScalarT = typename EvalT::ScalarT;

      /// Type for shared memory
      using scratch_view = Kokkos::View<ScalarT* ,typename PHX::DevLayout<ScalarT>::type,typename PHX::exec_space::scratch_memory_space,Kokkos::MemoryUnmanaged>;

      /**
       *  \brief An `enum` determining the behavior of this `Evaluator`.
       *
       *  This `Evaluator` will compute the result of its integration and then:
       *  - CONTRIBUTES:  contribute it to a specified residual, not saving
       *                  anything; or
       *  - EVALUATES:    save it under a specified name for future use.
       */
      const panzer::EvaluatorStyle evalStyle_;

      /**
       *  \brief A field to which we'll contribute, or in which we'll store,
       *         the result of computing this integral.
       */
      PHX::MDField<ScalarT, panzer::Cell, panzer::BASIS> field_;

      /**
       *  \brief A field representing the vector-valued function we're
       *         integrating (\f$ \vec{s} \f$) in a 2/3-dimensional problem.
       */
      PHX::MDField<const ScalarT, panzer::Cell, panzer::IP> vector2D_;
      PHX::MDField<const ScalarT, panzer::Cell, panzer::IP, panzer::Dim> vector3D_;

      /**
       *  \brief The (possibly empty) list of fields that are multipliers out
       *         in front of the integral (\f$ a(x) \f$, \f$ b(x) \f$, etc.).
       */
      PHX::MDField<const ScalarT, panzer::Cell, panzer::IP> Inv_Permeability_;

      /**
       *  \brief The `Kokkos::View` representation of the (possibly empty) list
       *         of fields that are multipliers out in front of the integral
       *         (\f$ a(x) \f$, \f$ b(x) \f$, etc.).
       */
      Kokkos::View<Kokkos::View<const ScalarT**,typename PHX::DevLayout<ScalarT>::type,PHX::Device>*> kokkosFieldMults_;

      /**
       *  \brief The name of the basis we're using.
       */
      std::string basisName_;

      /**
       *  \brief The index in the `Workset` bases for our particular
       *         `BasisIRLayout` name.
       */
      std::size_t basisIndex_;

      /**
       *  \brief The spatial dimension of the vector-valued function we're
       *         integrating, either 2 or 3.
       */
      int spaceDim_;

  }; // end of class Residual_Magnetostatics

}


#include "EquationSet_Magnetostatics_impl.hpp"

#endif
